import librosa
import numpy as np
from paddlespeech.server.utils.audio_process import change_speed
import soundfile as sf
import io
import base64

from paddlespeech.server.utils.exception import ServerBaseException


def postprocess(wav,
                original_fs: int,
                target_fs: int = 0,
                volume: float = 1.0,
                speed: float = 1.0,
                audio_path: str = None):
    """Post-processing operations, including speech, volume, sample rate, save audio file

    Args:
        wav (numpy(float)): Synthesized audio sample points
        original_fs (int): original audio sample rate
        target_fs (int): target audio sample rate
        volume (float): target volume
        speed (float): target speed

    Raises:
        ServerBaseException: Throws an exception if the change speed unsuccessfully.

    Returns:
        target_fs: target sample rate for synthesized audio.
        wav_base64: The base64 format of the synthesized audio.
    """

    # transform sample_rate
    if target_fs == 0 or target_fs > original_fs:
        target_fs = original_fs
        wav_tar_fs = wav
        print(
            "The sample rate of synthesized audio is the same as model, which is {}Hz".
            format(original_fs))
    else:
        wav_tar_fs = librosa.resample(
            np.squeeze(wav), original_fs, target_fs)
        print(
            "The sample rate of model is {}Hz and the target sample rate is {}Hz. Converting the sample rate of the synthesized audio successfully.".
            format(original_fs, target_fs))
    # transform volume
    wav_vol = wav_tar_fs * volume
    print("Transform the volume of the audio successfully.")

    # transform speed
    try:  # windows not support soxbindings
        wav_speed = change_speed(wav_vol, speed, target_fs)
        print("Transform the speed of the audio successfully.")
    except ServerBaseException:
        raise Exception(
            "Failed to transform speed. Can not install soxbindings on your system. \
             You need to set speed value 1.0.")
    except Exception as e:
        print("Failed to transform speed.")

    # wav to base64
    buf = io.BytesIO()
    sf.write(buf, wav_speed, target_fs, format="wav")
    buf.seek(0)
    base64_bytes = base64.b64encode(buf.read())
    wav_base64 = base64_bytes.decode('utf-8')
    print("Audio to string successfully.")

    # save audio
    if audio_path is not None:
        if audio_path.endswith(".wav"):
            sf.write(audio_path, wav_speed, target_fs)
        elif audio_path.endswith(".pcm"):
            wav_norm = wav_speed * (32767 / max(0.001,
                                                np.max(np.abs(wav_speed))))
            with open(audio_path, "wb") as f:
                f.write(wav_norm.astype(np.int16))
        print("Save audio to {} successfully.".format(audio_path))
    else:
        print("There is no need to save audio.")

    return target_fs, wav_base64


if __name__ == '__main__':
    y, sr = librosa.load('/home/output/test_e2e/001.wav')
    postprocess(y, sr, 0, speed=0.8, audio_path='/home/output/test_e2e/001_0.8x.wav')
